#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
流式ASR服务模块

基于火山引擎的WebSocket流式语音识别服务
支持实时音频流输入，边说边识别
"""

import asyncio
import base64
import gzip
import hmac
import json
import uuid
import wave
from hashlib import sha256
from io import BytesIO
from urllib.parse import urlparse
import websockets
from typing import Optional, Callable, List, Dict, Any
import pyaudio
import numpy as np
import threading
import time
import queue  # 使用标准库的线程安全队列

# 协议常量定义
PROTOCOL_VERSION = 0b0001
DEFAULT_HEADER_SIZE = 0b0001

# Message Type
CLIENT_FULL_REQUEST = 0b0001
CLIENT_AUDIO_ONLY_REQUEST = 0b0010
SERVER_FULL_RESPONSE = 0b1001
SERVER_ACK = 0b1011
SERVER_ERROR_RESPONSE = 0b1111

# Message Type Specific Flags
NO_SEQUENCE = 0b0000
POS_SEQUENCE = 0b0001
NEG_SEQUENCE = 0b0010
NEG_SEQUENCE_1 = 0b0011

# Message Serialization
NO_SERIALIZATION = 0b0000
JSON = 0b0001
THRIFT = 0b0011
CUSTOM_TYPE = 0b1111

# Message Compression
NO_COMPRESSION = 0b0000
GZIP = 0b0001
CUSTOM_COMPRESSION = 0b1111


def generate_header(
    version=PROTOCOL_VERSION,
    message_type=CLIENT_FULL_REQUEST,
    message_type_specific_flags=NO_SEQUENCE,
    serial_method=JSON,
    compression_type=GZIP,
    reserved_data=0x00,
    extension_header=bytes()
):
    """生成协议头"""
    header = bytearray()
    header_size = int(len(extension_header) / 4) + 1
    header.append((version << 4) | header_size)
    header.append((message_type << 4) | message_type_specific_flags)
    header.append((serial_method << 4) | compression_type)
    header.append(reserved_data)
    header.extend(extension_header)
    return header


def generate_full_default_header():
    """生成完整请求的默认头"""
    return generate_header()


def generate_audio_default_header():
    """生成音频数据的默认头"""
    return generate_header(message_type=CLIENT_AUDIO_ONLY_REQUEST)


def generate_last_audio_default_header():
    """生成最后一个音频包的默认头"""
    return generate_header(
        message_type=CLIENT_AUDIO_ONLY_REQUEST,
        message_type_specific_flags=NEG_SEQUENCE
    )


def parse_response(res):
    """解析服务器响应"""
    protocol_version = res[0] >> 4
    header_size = res[0] & 0x0f
    message_type = res[1] >> 4
    message_type_specific_flags = res[1] & 0x0f
    serialization_method = res[2] >> 4
    message_compression = res[2] & 0x0f
    reserved = res[3]
    header_extensions = res[4:header_size * 4]
    payload = res[header_size * 4:]
    
    result = {}
    payload_msg = None
    payload_size = 0
    
    if message_type == SERVER_FULL_RESPONSE:
        payload_size = int.from_bytes(payload[:4], "big", signed=True)
        payload_msg = payload[4:]
    elif message_type == SERVER_ACK:
        seq = int.from_bytes(payload[:4], "big", signed=True)
        result['seq'] = seq
        if len(payload) >= 8:
            payload_size = int.from_bytes(payload[4:8], "big", signed=False)
            payload_msg = payload[8:]
    elif message_type == SERVER_ERROR_RESPONSE:
        code = int.from_bytes(payload[:4], "big", signed=False)
        result['code'] = code
        payload_size = int.from_bytes(payload[4:8], "big", signed=False)
        payload_msg = payload[8:]
    
    if payload_msg is None:
        return result
    
    if message_compression == GZIP:
        payload_msg = gzip.decompress(payload_msg)
    
    if serialization_method == JSON:
        payload_msg = json.loads(str(payload_msg, "utf-8"))
    elif serialization_method != NO_SERIALIZATION:
        payload_msg = str(payload_msg, "utf-8")
    
    result['payload_msg'] = payload_msg
    result['payload_size'] = payload_size
    return result


class StreamingASRService:
    """
    流式ASR服务
    
    实现实时音频流的语音识别，边说边识别
    支持异步处理和回调机制
    """
    
    def __init__(self, 
                 appid: str = "1871057448",
                 token: str = "tnph6TBEuRv4MWgh425-XUf6K7AcS2_5",
                 cluster: str = "volcengine_streaming_common",
                 ws_url: str = "wss://openspeech.bytedance.com/api/v2/asr",
                 **kwargs):
        """
        初始化流式ASR服务
        
        Args:
            appid: 火山引擎应用ID
            token: 访问令牌
            cluster: 集群名称
            ws_url: WebSocket服务地址
            **kwargs: 其他配置参数
        """
        self.appid = appid
        self.token = token
        self.cluster = cluster
        self.ws_url = ws_url
        
        # 音频参数
        self.sample_rate = kwargs.get("sample_rate", 16000)
        self.channels = kwargs.get("channels", 1)
        self.bits = kwargs.get("bits", 16)
        self.format = kwargs.get("format", "raw")
        self.codec = kwargs.get("codec", "raw")
        self.language = kwargs.get("language", "zh-CN")
        
        # 识别参数
        self.nbest = kwargs.get("nbest", 1)
        self.show_utterances = kwargs.get("show_utterances", True)
        self.result_type = kwargs.get("result_type", "full")
        self.workflow = kwargs.get("workflow", "audio_in,resample,partition,vad,fe,decode,itn,nlu_punctuate")
        
        # 用户标识
        self.uid = kwargs.get("uid", "streaming_asr_user")
        
        # 状态
        self.success_code = 1000
        self.is_running = False
        self.ws = None
        
        # 回调函数
        self.on_result: Optional[Callable] = kwargs.get('on_result', None)  # 识别结果回调（中间+最终）
        self.on_final: Optional[Callable] = kwargs.get('on_final', None)    # 最终结果回调
        self.on_error: Optional[Callable] = kwargs.get('on_error', None)    # 错误回调
        self.on_complete: Optional[Callable] = kwargs.get('on_complete', None)  # 完成回调
        
        # 去重缓存：防止相同的最终结果被重复触发
        self.last_final_text = ""
        self.last_partial_text = ""
        
        print(f"🔊 流式ASR服务初始化完成")
        print(f"   AppID: {appid}")
        print(f"   集群: {cluster}")
        print(f"   采样率: {self.sample_rate}Hz")
    
    def construct_request(self, reqid: str) -> dict:
        """构建识别请求"""
        req = {
            'app': {
                'appid': self.appid,
                'cluster': self.cluster,
                'token': self.token,
            },
            'user': {
                'uid': self.uid
            },
            'request': {
                'reqid': reqid,
                'nbest': self.nbest,
                'workflow': self.workflow,
                'show_utterances': self.show_utterances,
                'result_type': self.result_type,
                "sequence": 1
            },
            'audio': {
                'format': self.format,
                'rate': self.sample_rate,
                'language': self.language,
                'bits': self.bits,
                'channel': self.channels,
                'codec': self.codec
            }
        }
        return req
    
    def token_auth(self) -> dict:
        """Token认证"""
        return {'Authorization': 'Bearer; {}'.format(self.token)}
    
    async def connect(self):
        """建立WebSocket连接"""
        try:
            header = self.token_auth()
            self.ws = await websockets.connect(
                self.ws_url, 
                additional_headers=header, 
                max_size=1000000000
            )
            
            # 发送初始化请求
            reqid = str(uuid.uuid4())
            request_params = self.construct_request(reqid)
            payload_bytes = str.encode(json.dumps(request_params))
            payload_bytes = gzip.compress(payload_bytes)
            
            full_client_request = bytearray(generate_full_default_header())
            full_client_request.extend((len(payload_bytes)).to_bytes(4, 'big'))
            full_client_request.extend(payload_bytes)
            
            await self.ws.send(full_client_request)
            
            # 接收初始化响应
            res = await self.ws.recv()
            result = parse_response(res)
            
            if 'payload_msg' in result:
                if result['payload_msg'].get('code') != self.success_code:
                    error_msg = result['payload_msg'].get('message', '未知错误')
                    print(f"❌ 连接失败: {error_msg}")
                    if self.on_error:
                        self.on_error(error_msg)
                    return False
            
            self.is_running = True
            print("✅ WebSocket连接建立成功")
            return True
            
        except Exception as e:
            print(f"❌ 连接异常: {e}")
            if self.on_error:
                self.on_error(str(e))
            return False
    
    async def send_audio(self, audio_data: bytes, is_last: bool = False):
        """
        发送音频数据（仅发送，不接收）
        
        Args:
            audio_data: 音频数据（PCM格式）
            is_last: 是否是最后一帧
        """
        if not self.ws or not self.is_running:
            return
        
        try:
            # 压缩音频数据
            payload_bytes = gzip.compress(audio_data)
            
            # 选择合适的header
            if is_last:
                audio_request = bytearray(generate_last_audio_default_header())
            else:
                audio_request = bytearray(generate_audio_default_header())
            
            audio_request.extend((len(payload_bytes)).to_bytes(4, 'big'))
            audio_request.extend(payload_bytes)
            
            # 仅发送音频数据，不接收
            await self.ws.send(audio_request)
            
        except Exception as e:
            print(f"❌ 发送音频异常: {e}")
            if self.on_error:
                self.on_error(str(e))
    
    async def receive_results(self):
        """
        持续接收识别结果（单独的协程）
        """
        result_count = 0
        try:
            print("🎧 开始接收识别结果...")
            while self.is_running and self.ws:
                try:
                    # 接收识别结果
                    res = await self.ws.recv()
                    result_count += 1
                    print(f"📥 收到第 {result_count} 个响应包")
                    
                    result = parse_response(res)
                    print(f"🔍 解析结果: {result}")
                    
                    # 处理识别结果
                    if 'payload_msg' in result:
                        payload = result['payload_msg']
                        print(f"📦 payload内容: {payload}")
                        
                        if payload.get('code') != self.success_code:
                            error_msg = payload.get('message', '识别错误')
                            error_code = payload.get('code')
                            backend_code = payload.get('backend_code')
                            
                            print(f"⚠️ 识别错误 (code={error_code}, backend_code={backend_code}): {error_msg}")
                            
                            # 检查是否是会话超时或已结束的错误
                            if backend_code in [1020, 45000000]:
                                # 1020: 超时错误
                                # 45000000: 会话已结束
                                print(f"🔴 检测到会话失效错误，停止接收并断开连接")
                                self.is_running = False
                                if self.on_error:
                                    self.on_error(f"会话已失效: {error_msg}")
                                # 立即断开WebSocket连接，确保资源清理
                                try:
                                    if self.ws:
                                        await self.ws.close()
                                        self.ws = None
                                        print("🔌 已主动关闭失效的WebSocket连接")
                                except Exception as close_error:
                                    print(f"⚠️ 关闭连接时出错: {close_error}")
                                break  # 退出接收循环
                            else:
                                # 其他错误，继续接收
                                if self.on_error:
                                    self.on_error(error_msg)
                        else:
                            # 提取识别文本（result是一个列表）
                            result_list = payload.get('result', [])
                            text = ''
                            is_final = False
                            
                            # 从结果列表中提取文本
                            if result_list and isinstance(result_list, list) and len(result_list) > 0:
                                first_result = result_list[0]
                                
                                # 🔍 调试：打印完整结果结构（仅前3次）
                                if not hasattr(self, '_debug_count'):
                                    self._debug_count = 0
                                if self._debug_count < 3:
                                    import json
                                    print(f"🔍 [调试 {self._debug_count + 1}/3] ASR返回结构:")
                                    print(json.dumps(first_result, ensure_ascii=False, indent=2)[:1000])
                                    self._debug_count += 1
                                
                                text = first_result.get('text', '')
                                
                                # 检查是否是最终结果
                                # 策略：只有当所有utterances都是definite=True时，才认为是最终结果
                                utterances = first_result.get('utterances', [])
                                if utterances:
                                    # 检查是否所有utterances都已确定
                                    all_definite = all(
                                        utt.get('definite', False) or not utt.get('text', '').strip()
                                        for utt in utterances
                                    )
                                    # 如果所有非空utterances都是definite=True，认为是最终结果
                                    has_definite = any(
                                        utt.get('definite', False) and utt.get('text', '').strip()
                                        for utt in utterances
                                    )
                                    
                                    # 只有当有确定的内容且所有内容都已确定时，才是最终结果
                                    is_final = has_definite and all_definite
                            
                            print(f"📝 识别文本: '{text}', is_final: {is_final}")
                            
                            if text:
                                # 去重检查：防止相同结果被重复触发
                                if is_final:
                                    # 最终结果去重
                                    if text == self.last_final_text:
                                        print(f"⚠️ [ASR去重] 跳过重复的最终结果: {text[:30]}...")
                                        continue  # 跳过重复的最终结果
                                    else:
                                        print(f"✅ [ASR去重] 新的最终结果: {text[:30]}...")
                                        self.last_final_text = text
                                        self.last_partial_text = ""  # 清空中间结果缓存
                                else:
                                    # 中间结果去重
                                    if text == self.last_partial_text:
                                        # print(f"⚠️ [ASR去重] 跳过重复的中间结果")
                                        continue  # 跳过重复的中间结果
                                    else:
                                        self.last_partial_text = text
                                
                                # 触发回调
                                if self.on_result:
                                    result_data = {
                                        'text': text,
                                        'is_final': is_final,
                                        'payload': payload
                                    }
                                    self.on_result(result_data)
                                
                                # 如果是最终结果，也调用on_final
                                if is_final and self.on_final:
                                    self.on_final(result_data)
                                
                                # 调用完成回调
                                if is_final and self.on_complete:
                                    self.on_complete(payload)
                            else:
                                print("⚠️ 识别文本为空")
                    else:
                        print(f"⚠️ 响应中没有payload_msg: {result}")
                
                except Exception as e:
                    if self.is_running:  # 只有在运行中才报错
                        print(f"❌ 接收结果异常: {e}")
                        import traceback
                        traceback.print_exc()
                    break
                    
        except Exception as e:
            print(f"❌ 接收循环异常: {e}")
            import traceback
            traceback.print_exc()
        finally:
            print(f"🛑 接收循环结束，共收到 {result_count} 个响应")
    
    async def disconnect(self):
        """断开WebSocket连接"""
        self.is_running = False
        if self.ws:
            await self.ws.close()
            self.ws = None
            print("🔌 WebSocket连接已关闭")
        
        # 清理去重缓存
        self.last_final_text = ""
        self.last_partial_text = ""
        print("🧹 已清理ASR去重缓存")


class RealTimeStreamingASR:
    """
    实时流式ASR服务
    
    结合麦克风实时采集和流式ASR识别
    实现边说边识别的功能
    """
    
    def __init__(self, 
                 appid: str = "1871057448",
                 token: str = "tnph6TBEuRv4MWgh425-XUf6K7AcS2_5",
                 cluster: str = "volcengine_streaming_common",
                 sample_rate: int = 16000,
                 channels: int = 1,
                 chunk_size: int = 3200):  # 每次发送0.2秒的音频
        """
        初始化实时流式ASR
        
        Args:
            appid: 应用ID
            token: 访问令牌
            cluster: 集群名称
            sample_rate: 采样率
            channels: 声道数
            chunk_size: 音频块大小（字节）
        """
        # 创建流式ASR服务
        self.asr_service = StreamingASRService(
            appid=appid,
            token=token,
            cluster=cluster,
            sample_rate=sample_rate,
            channels=channels
        )
        
        # 音频参数
        self.sample_rate = sample_rate
        self.channels = channels
        self.chunk_size = chunk_size
        self.audio_format = pyaudio.paInt16
        
        # PyAudio实例
        self.audio = pyaudio.PyAudio()
        self.stream = None
        
        # 状态控制
        self.is_listening = False
        self.audio_queue = queue.Queue(maxsize=100)  # 使用线程安全队列
        
        # 设置回调
        self.asr_service.on_result = self._on_result
        self.asr_service.on_error = self._on_error
        self.asr_service.on_complete = self._on_complete
        
        # 用户自定义回调
        self.on_partial_result: Optional[Callable] = None  # 中间结果
        self.on_final_result: Optional[Callable] = None    # 最终结果
        
        print(f"🎤 实时流式ASR初始化完成")
    
    def _on_result(self, result_data: dict):
        """识别结果回调"""
        text = result_data.get('text', '')
        is_final = result_data.get('is_final', False)
        
        if is_final:
            print(f"✅ [最终] {text}")
            if self.on_final_result:
                self.on_final_result(text, result_data)
        else:
            print(f"💬 [中间] {text}")
            if self.on_partial_result:
                self.on_partial_result(text, result_data)
    
    def _on_error(self, error_msg: str):
        """错误回调"""
        print(f"❌ 错误: {error_msg}")
    
    def _on_complete(self, payload: dict):
        """完成回调"""
        print(f"🎉 识别完成")
    
    def list_audio_devices(self):
        """列出可用的音频设备"""
        print("🎧 可用音频设备:")
        for i in range(self.audio.get_device_count()):
            device_info = self.audio.get_device_info_by_index(i)
            if device_info['maxInputChannels'] > 0:
                print(f"  [{i}] {device_info['name']} (输入声道: {device_info['maxInputChannels']})")
    
    async def _audio_capture_loop(self):
        """音频采集循环"""
        print("🎤 开始音频采集...")
        
        def audio_callback(in_data, frame_count, time_info, status):
            """音频流回调"""
            if self.is_listening:
                # 将音频数据放入线程安全队列
                try:
                    self.audio_queue.put_nowait(in_data)
                except queue.Full:
                    # 队列满了，丢弃旧数据
                    pass
            return (in_data, pyaudio.paContinue)
        
        # 打开音频流
        self.stream = self.audio.open(
            format=self.audio_format,
            channels=self.channels,
            rate=self.sample_rate,
            input=True,
            frames_per_buffer=self.chunk_size,
            stream_callback=audio_callback
        )
        
        self.stream.start_stream()
        print("✅ 音频流已启动")
    
    async def _audio_send_loop(self):
        """音频发送循环"""
        print("📤 开始音频发送...")
        
        while self.is_listening:
            try:
                # 从线程安全队列获取音频数据（非阻塞）
                try:
                    audio_data = self.audio_queue.get(timeout=0.1)
                except queue.Empty:
                    # 队列为空，继续等待
                    await asyncio.sleep(0.01)
                    continue
                
                # 发送音频数据到ASR服务
                await self.asr_service.send_audio(audio_data, is_last=False)
                
            except Exception as e:
                print(f"❌ 发送音频异常: {e}")
                break
        
        # 发送最后一帧
        if not self.audio_queue.empty():
            try:
                audio_data = self.audio_queue.get_nowait()
                await self.asr_service.send_audio(audio_data, is_last=True)
            except queue.Empty:
                # 发送一个空帧标记结束
                await self.asr_service.send_audio(b'', is_last=True)
        else:
            # 发送一个空帧标记结束
            await self.asr_service.send_audio(b'', is_last=True)
        
        print("🔚 音频发送结束")
    
    async def start(self, device_index: Optional[int] = None):
        """
        开始实时识别
        
        Args:
            device_index: 音频设备索引，None为默认设备
        """
        if self.is_listening:
            print("⚠️ 已在监听中")
            return
        
        print("🚀 启动实时流式ASR...")
        
        # 连接到ASR服务
        if not await self.asr_service.connect():
            print("❌ 连接ASR服务失败")
            return
        
        self.is_listening = True
        
        # 启动音频采集和发送
        capture_task = asyncio.create_task(self._audio_capture_loop())
        send_task = asyncio.create_task(self._audio_send_loop())
        
        print("✅ 实时识别已启动，开始说话...")
        print("   按 Ctrl+C 停止")
        
        try:
            # 等待任务完成
            await asyncio.gather(capture_task, send_task)
        except asyncio.CancelledError:
            print("\n🛑 任务已取消")
    
    async def stop(self):
        """停止实时识别"""
        print("\n🛑 停止实时识别...")
        
        self.is_listening = False
        
        # 停止音频流
        if self.stream:
            self.stream.stop_stream()
            self.stream.close()
            self.stream = None
        
        # 断开ASR连接
        await self.asr_service.disconnect()
        
        print("✅ 已停止")
    
    def __del__(self):
        """析构函数"""
        if hasattr(self, 'audio'):
            self.audio.terminate()


async def main():
    """主函数 - 测试实时流式ASR"""
    print("=" * 60)
    print("实时流式ASR测试")
    print("=" * 60)
    
    # 创建实时流式ASR实例
    asr = RealTimeStreamingASR(
        appid="1871057448",
        token="tnph6TBEuRv4MWgh425-XUf6K7AcS2_5",
        cluster="volcengine_streaming_common",
        sample_rate=16000,
        channels=1
    )
    
    # 列出音频设备
    asr.list_audio_devices()
    
    # 设置结果回调
    def on_partial(text, data):
        """中间结果回调"""
        pass  # 已在默认回调中打印
    
    def on_final(text, data):
        """最终结果回调"""
        pass  # 已在默认回调中打印
    
    asr.on_partial_result = on_partial
    asr.on_final_result = on_final
    
    try:
        # 开始实时识别
        await asr.start()
    except KeyboardInterrupt:
        print("\n⚠️ 用户中断")
    finally:
        await asr.stop()
        print("\n👋 测试结束")


if __name__ == "__main__":
    asyncio.run(main())

